import pandas as pd
import numpy as np
from typing import List, Dict, Optional
from decimal import Decimal
from django.db import models
from stocks.models import DailyData


class TechnicalAnalysis:
    """技术指标计算类"""
    
    @staticmethod
    def calculate_ma(data: List[Dict], period: int) -> List[Optional[float]]:
        """
        计算移动平均线
        
        Args:
            data: 股票数据列表，每个元素包含'close'价格
            period: 移动平均周期
            
        Returns:
            移动平均值列表，长度与输入数据相同，前面period-1个为None
        """
        if not data or len(data) < period:
            return [None] * len(data)
            
        closes = [float(item['close']) if item['close'] is not None else None for item in data]
        ma_values = []
        
        for i in range(len(closes)):
            if i < period - 1:
                ma_values.append(None)
            else:
                window = closes[i - period + 1: i + 1]
                # 过滤掉None值
                valid_values = [v for v in window if v is not None]
                if valid_values:
                    ma_values.append(sum(valid_values) / len(valid_values))
                else:
                    ma_values.append(None)
        
        return ma_values
    
    @staticmethod
    def calculate_rsi(data: List[Dict], period: int = 14) -> List[Optional[float]]:
        """
        计算RSI相对强弱指标
        
        Args:
            data: 股票数据列表
            period: RSI计算周期，默认14
            
        Returns:
            RSI值列表
        """
        if len(data) < period + 1:
            return [None] * len(data)
            
        closes = [float(item['close']) if item['close'] is not None else None for item in data]
        changes = []
        
        # 计算价格变化
        for i in range(1, len(closes)):
            if closes[i] is not None and closes[i-1] is not None:
                changes.append(closes[i] - closes[i-1])
            else:
                changes.append(None)
        
        rsi_values = [None] * len(data)
        
        for i in range(period, len(changes)):
            window = changes[i - period: i]
            gains = [max(change, 0) for change in window if change is not None]
            losses = [max(-change, 0) for change in window if change is not None]
            
            if not gains or not losses:
                continue
                
            avg_gain = sum(gains) / period
            avg_loss = sum(losses) / period
            
            if avg_loss == 0:
                rsi = 100
            else:
                rs = avg_gain / avg_loss
                rsi = 100 - (100 / (1 + rs))
            
            rsi_values[i + 1] = rsi  # i+1因为changes比closes少一个元素
        
        return rsi_values
    
    @staticmethod
    def calculate_macd(data: List[Dict], fast_period: int = 12, slow_period: int = 26, signal_period: int = 9) -> Dict:
        """
        计算MACD指标
        
        Returns:
            {
                'macd': MACD值列表,
                'signal': 信号线列表,
                'histogram': 柱状图列表
            }
        """
        if len(data) < slow_period + signal_period:
            empty = [None] * len(data)
            return {'macd': empty, 'signal': empty, 'histogram': empty}
            
        closes = [float(item['close']) if item['close'] is not None else None for item in data]
        
        # 计算EMA
        def calculate_ema(prices, period):
            if not prices or len(prices) < period:
                return [None] * len(prices)
                
            ema_values = [None] * len(prices)
            multiplier = 2 / (period + 1)
            
            # 第一个EMA是简单移动平均
            valid_prices = [p for p in prices[:period] if p is not None]
            if valid_prices:
                ema_values[period - 1] = sum(valid_prices) / len(valid_prices)
            
            for i in range(period, len(prices)):
                if prices[i] is not None and ema_values[i-1] is not None:
                    ema_values[i] = (prices[i] - ema_values[i-1]) * multiplier + ema_values[i-1]
                else:
                    ema_values[i] = ema_values[i-1] if i > 0 else None
            
            return ema_values
        
        ema_fast = calculate_ema(closes, fast_period)
        ema_slow = calculate_ema(closes, slow_period)
        
        # 计算MACD线
        macd_line = []
        for i in range(len(closes)):
            if ema_fast[i] is not None and ema_slow[i] is not None:
                macd_line.append(ema_fast[i] - ema_slow[i])
            else:
                macd_line.append(None)
        
        # 计算信号线（MACD的EMA）
        signal_line = calculate_ema(macd_line, signal_period)
        
        # 计算柱状图
        histogram = []
        for i in range(len(closes)):
            if macd_line[i] is not None and signal_line[i] is not None:
                histogram.append(macd_line[i] - signal_line[i])
            else:
                histogram.append(None)
        
        return {
            'macd': macd_line,
            'signal': signal_line,
            'histogram': histogram
        }
    
    @staticmethod
    def calculate_bollinger_bands(data: List[Dict], period: int = 20, std_dev: int = 2) -> Dict:
        """
        计算布林带
        
        Returns:
            {
                'middle': 中轨（移动平均）,
                'upper': 上轨,
                'lower': 下轨
            }
        """
        if len(data) < period:
            empty = [None] * len(data)
            return {'middle': empty, 'upper': empty, 'lower': empty}
            
        closes = [float(item['close']) if item['close'] is not None else None for item in data]
        
        middle_band = TechnicalAnalysis.calculate_ma(data, period)
        upper_band = [None] * len(data)
        lower_band = [None] * len(data)
        
        for i in range(period - 1, len(closes)):
            if middle_band[i] is not None:
                window = closes[i - period + 1: i + 1]
                valid_values = [v for v in window if v is not None]
                
                if valid_values:
                    std = np.std(valid_values)
                    upper_band[i] = middle_band[i] + std_dev * std
                    lower_band[i] = middle_band[i] - std_dev * std
        
        return {
            'middle': middle_band,
            'upper': upper_band,
            'lower': lower_band
        }


class StockDataProcessor:
    """股票数据处理工具类"""
    
    @staticmethod
    def get_stock_data(symbol: str, days: int = 100) -> List[Dict]:
        """
        获取股票历史数据
        
        Args:
            symbol: 股票代码
            days: 获取最近多少天的数据
            
        Returns:
            股票数据列表，按日期倒序排列
        """
        from django.db import connections
        
        # 使用stock_daily_view视图查询数据
        sql_query = """
            SELECT id, symbol, trade_date, open, close, high, low, volume, amount, 
                   price_amplitude, change_pct, change_amount, turnover_rate
            FROM stock_daily_view
            WHERE symbol = %s
            ORDER BY trade_date DESC 
            LIMIT %s
        """
        
        data = []
        with connections['postgres'].cursor() as cursor:
            try:
                cursor.execute(sql_query, [symbol, days])
                if cursor.description:
                    columns = [col[0] for col in cursor.description]
                    for row in cursor.fetchall():
                        row_dict = dict(zip(columns, row))
                        data.append({
                            'date': row_dict['trade_date'],
                            'open': float(row_dict['open']) if row_dict['open'] else None,
                            'close': float(row_dict['close']) if row_dict['close'] else None,
                            'high': float(row_dict['high']) if row_dict['high'] else None,
                            'low': float(row_dict['low']) if row_dict['low'] else None,
                            'volume': row_dict['volume'],
                            'amount': float(row_dict['amount']) if row_dict['amount'] else None
                        })
            except Exception as e:
                print(f"查询股票数据时出错: {e}")
                return []
        
        # 反转数据，使日期正序排列
        return list(reversed(data))
    
    @staticmethod
    def calculate_kdj(data: List[Dict], period: int = 9, k_period: int = 3, d_period: int = 3) -> Dict:
        """
        计算KDJ指标
        
        Returns:
            {
                'k': K值列表,
                'd': D值列表,
                'j': J值列表
            }
        """
        if len(data) < period:
            empty = [None] * len(data)
            return {'k': empty, 'd': empty, 'j': empty}
            
        highs = [float(item['high']) if item['high'] is not None else None for item in data]
        lows = [float(item['low']) if item['low'] is not None else None for item in data]
        closes = [float(item['close']) if item['close'] is not None else None for item in data]
        
        k_values = [None] * len(data)
        d_values = [None] * len(data)
        j_values = [None] * len(data)
        
        for i in range(period - 1, len(data)):
            # 计算最近period天内的最高价和最低价
            window_highs = highs[i - period + 1: i + 1]
            window_lows = lows[i - period + 1: i + 1]
            
            valid_highs = [h for h in window_highs if h is not None]
            valid_lows = [l for l in window_lows if l is not None]
            
            if not valid_highs or not valid_lows or closes[i] is None:
                continue
                
            highest_high = max(valid_highs)
            lowest_low = min(valid_lows)
            
            if highest_high == lowest_low:
                rsv = 50  # 避免除零错误
            else:
                rsv = (closes[i] - lowest_low) / (highest_high - lowest_low) * 100
            
            # 计算K值（RSV的移动平均）
            if k_values[i-1] is None:
                k_values[i] = rsv
            else:
                k_values[i] = (2/3) * k_values[i-1] + (1/3) * rsv
            
            # 计算D值（K值的移动平均）
            if d_values[i-1] is None:
                d_values[i] = k_values[i]
            else:
                d_values[i] = (2/3) * d_values[i-1] + (1/3) * k_values[i]
            
            # 计算J值
            if k_values[i] is not None and d_values[i] is not None:
                j_values[i] = 3 * k_values[i] - 2 * d_values[i]
        
        return {'k': k_values, 'd': d_values, 'j': j_values}

    @staticmethod
    def prepare_analysis_data(symbol: str, days: int = 100) -> Dict:
        """
        准备技术分析所需的数据
        
        Returns:
            包含原始数据和技术指标的数据字典
        """
        raw_data = StockDataProcessor.get_stock_data(symbol, days)
        
        if not raw_data:
            return {'raw_data': [], 'indicators': {}}
        
        # 计算各种技术指标
        ma5 = TechnicalAnalysis.calculate_ma(raw_data, 5)
        ma10 = TechnicalAnalysis.calculate_ma(raw_data, 10)
        ma20 = TechnicalAnalysis.calculate_ma(raw_data, 20)
        ma60 = TechnicalAnalysis.calculate_ma(raw_data, 60)
        
        rsi = TechnicalAnalysis.calculate_rsi(raw_data, 14)
        
        macd_result = TechnicalAnalysis.calculate_macd(raw_data)
        
        bollinger_result = TechnicalAnalysis.calculate_bollinger_bands(raw_data)
        
        kdj_result = StockDataProcessor.calculate_kdj(raw_data)
        
        return {
            'raw_data': raw_data,
            'indicators': {
                'ma5': ma5,
                'ma10': ma10,
                'ma20': ma20,
                'ma60': ma60,
                'rsi': rsi,
                'macd': macd_result['macd'],
                'signal': macd_result['signal'],
                'histogram': macd_result['histogram'],
                'bollinger_middle': bollinger_result['middle'],
                'bollinger_upper': bollinger_result['upper'],
                'bollinger_lower': bollinger_result['lower'],
                'k': kdj_result['k'],
                'd': kdj_result['d'],
                'j': kdj_result['j']
            }
        }
